#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
使用梯形加速度规划直线轨迹并可视化
"""

import numpy as np
from diff_robot_model import build_robot_model
from trajectory_planner import TrapezoidalTrajectoryPlanner, simulate_robot_trajectory
from draw import Draw_MPC_point_stabilization_v1

def main():
    # 创建差分小车模型
    robot_model = build_robot_model(
        T=0.1,  # 采样时间
        N=100,   # 预测步长
        rob_diam=0.3,  # 机器人直径
        a_max=1.0,     # 最大加速度
        omega_max=np.pi/4  # 最大角速度
    )
    
    print("差分小车模型参数:")
    print(f"  - 采样时间: {robot_model.T}s")
    print(f"  - 预测步长: {robot_model.N}")
    print(f"  - 机器人直径: {robot_model.rob_diam}m")
    print(f"  - 最大加速度: {robot_model.a_max}m/s²")
    print(f"  - 最大角速度: {robot_model.omega_max}rad/s")
    
    # 创建轨迹规划器
    planner = TrapezoidalTrajectoryPlanner(
        max_acceleration=robot_model.a_max,
        max_velocity=0.5  # 设置最大速度
    )
    
    # 定义起始点和终点
    start_pos = [0.0, 0.0]
    end_pos = [2.0, 0.0]
    duration = 5.0  # 5秒内完成轨迹
    
    print(f"\n轨迹规划参数:")
    print(f"  - 起始位置: {start_pos}")
    print(f"  - 终点位置: {end_pos}")
    print(f"  - 总时间: {duration}s")
    
    # 规划直线轨迹
    time_points, positions, velocities, accelerations = planner.plan_linear_trajectory(
        start_pos, end_pos, duration
    )
    
    print(f"\n轨迹规划完成:")
    print(f"  - 时间点数量: {len(time_points)}")
    print(f"  - 总距离: {np.linalg.norm(np.array(end_pos) - np.array(start_pos)):.2f}m")
    
    # 使用差分小车模型模拟轨迹
    robot_states, robot_controls = simulate_robot_trajectory(
        robot_model, positions, velocities, time_points
    )
    
    print(f"\n机器人模拟完成:")
    print(f"  - 状态数量: {len(robot_states)}")
    print(f"  - 控制数量: {len(robot_controls)}")
    
    # 提取用于可视化的状态 (x, y, theta)
    visualization_states = []
    for state in robot_states:
        visualization_states.append([state[0], state[1], state[4]])  # [x, y, theta]
    
    # 获取初始状态和目标状态
    init_state = np.array([start_pos[0], start_pos[1], 0.0])  # [x, y, theta]
    target_state = np.array([end_pos[0], end_pos[1], np.arctan2(velocities[-1][1], velocities[-1][0])])  # [x, y, theta]
    
    print(f"\n可视化参数:")
    print(f"  - 初始状态: {init_state}")
    print(f"  - 目标状态: {target_state}")
    
    # 可视化轨迹
    print(f"\n开始可视化...")
    try:
        draw_result = Draw_MPC_point_stabilization_v1(
            robot_states=visualization_states,
            init_state=init_state,
            target_state=target_state,
            rob_diam=robot_model.rob_diam,
            export_fig=False
        )
        print("可视化完成!")
    except Exception as e:
        print(f"可视化失败: {e}")
        # 打印一些关键信息用于调试
        print("可视化调试信息:")
        print(f"  - robot_states长度: {len(visualization_states)}")
        if len(visualization_states) > 0:
            print(f"  - 第一个状态: {visualization_states[0]}")
            print(f"  - 最后一个状态: {visualization_states[-1]}")
    
    # 打印轨迹统计信息
    print(f"\n轨迹统计:")
    final_position = [robot_states[-1][0], robot_states[-1][1]]
    position_error = np.linalg.norm(np.array(final_position) - np.array(end_pos))
    print(f"  - 最终位置: [{final_position[0]:.3f}, {final_position[1]:.3f}]")
    print(f"  - 目标位置: {end_pos}")
    print(f"  - 位置误差: {position_error:.3f}m")
    
    # 计算最大速度和加速度
    speeds = [np.linalg.norm([state[2], state[3]]) for state in robot_states]
    max_speed = max(speeds) if speeds else 0
    print(f"  - 最大速度: {max_speed:.3f}m/s")
    print(f"  - 速度设定上限: {planner.max_velocity}m/s")

if __name__ == "__main__":
    main()